# ====================================================================================
# === Louvain H3-Boundary Community Classifier (v5.0) ================================
# ====================================================================================
# 描述:
#   本脚本使用Louvain算法发现路网社区，并使用H3六边形生成精确的“锯齿状”边界，
#   最后对这些社区进行功能分类。
#   1. 使用Louvain算法发现路网社区。
#   2. [核心] 使用h3.h3_line()为社区描绘精确的、锯齿状的边界。
#   3. 为每个社区计算特征。
#   4. 使用K-Means对社区进行功能分类并赋予颜色。
#
# 作者: Gemini
# 日期: 2025-10-13
# ====================================================================================

import warnings
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point, Polygon, LineString
import networkx as nx
import community as community_louvain
import momepy
import folium
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
import h3
from collections import defaultdict, Counter

# --- 0. 环境与参数配置 ---
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=UserWarning)

# --- 文件路径定义 ---
BOUNDARY_SHP = "仓山区最终完整边界.shp"
ROAD_SHP = "仓山区修复后路网.shp"
POI_CSV = "poi仓山区.csv" 
STORES_CSV = "福州市仓山区瑞幸咖啡坐标.csv"

# --- 模型参数 ---
LOUVAIN_RESOLUTION = 0.5 
H3_RESOLUTION = 10 # H3描边分辨率，可以高一些以获得更精确的边界
MIN_COMMUNITY_SIZE = 10  
NUM_CLUSTERS = 4 # 功能区数量

# --- 输出文件 ---
OUTPUT_MAP_HTML = "community_h3_classification_map.html" 

def load_data(boundary_path, road_path, poi_path, stores_path):
    # ... (此函数无需修改) ...
    print("--- 阶段1: 正在加载并预处理数据 ---")
    try:
        boundary_gdf = gpd.read_file(boundary_path)
        if boundary_gdf.crs is None: boundary_gdf.set_crs("EPSG:4326", inplace=True)
        boundary_gdf_wgs84 = boundary_gdf.to_crs("EPSG:4326")
        road_gdf = gpd.read_file(road_path)
        if road_gdf.crs is None: road_gdf.set_crs("EPSG:4326", inplace=True)
        road_gdf_wgs84 = road_gdf.to_crs("EPSG:4326")
        poi_df = pd.read_csv(poi_path, encoding='utf-8')
        if 'category' not in poi_df.columns: raise ValueError("POI CSV文件中缺少 'category' 列！")
        poi_gdf = gpd.GeoDataFrame(poi_df, geometry=gpd.points_from_xy(poi_df.geolocation_lng, poi_df.geolocation_lat), crs="EPSG:4326")
        stores_df = pd.read_csv(stores_path, encoding='utf-8')
        stores_gdf = gpd.GeoDataFrame(stores_df, geometry=gpd.points_from_xy(stores_df.geolocation_lng, stores_df.geolocation_lat), crs="EPSG:4326")
        road_gdf_clipped = gpd.clip(road_gdf_wgs84, boundary_gdf_wgs84)
        poi_gdf_clipped = gpd.clip(poi_gdf, boundary_gdf_wgs84)
        print("✅ 数据加载和裁剪成功。")
        return road_gdf_clipped, poi_gdf_clipped, stores_gdf, boundary_gdf_wgs84
    except Exception as e:
        print(f"❌ 加载数据时出错: {e}")
        return None, None, None, None

def perform_louvain_community_detection(road_gdf, resolution_param):
    # ... (此函数无需修改) ...
    print("\n--- 阶段2: 正在执行路网社区发现 (Louvain) ---")
    road_gdf_exploded = road_gdf.explode(index_parts=True)
    G = momepy.gdf_to_nx(road_gdf_exploded, approach="primal")
    print(f"  - 正在使用 Louvain resolution = {resolution_param} 进行计算...")
    partition = community_louvain.best_partition(G, resolution=resolution_param)
    nx.set_node_attributes(G, partition, 'community')
    num_communities = len(set(partition.values()))
    print(f"✅ Louvain算法完成，初步发现了 {num_communities} 个社区。")
    return G

# === [核心函数: H3描边法] ===
def delineate_communities_with_h3(G, resolution, min_node_count):
    print(f"\n--- 阶段3: 正在使用 H3 (分辨率: {resolution}) 描绘精确边界 ---")
    community_counts = Counter(nx.get_node_attributes(G, 'community').values())
    hex_community_votes = defaultdict(list)
    
    for u, v in G.edges():
        u_comm = G.nodes[u].get('community')
        if community_counts.get(u_comm, 0) < min_node_count: continue
        u_node, v_node = G.nodes[u], G.nodes[v]
        if 'x' not in u_node or 'y' not in u_node or 'x' not in v_node or 'y' not in v_node: continue
        line = LineString([(u_node['x'], u_node['y']), (v_node['x'], v_node['y'])])
        coords = list(line.coords)
        if len(coords) < 2: continue
        edge_hexagons = set()
        for i in range(len(coords) - 1):
            start_lon, start_lat = coords[i]
            end_lon, end_lat = coords[i+1]
            start_hex = h3.geo_to_h3(start_lat, start_lon, resolution)
            end_hex = h3.geo_to_h3(end_lat, end_lon, resolution)
            line_path_hexs = set(h3.h3_line(start_hex, end_hex))
            edge_hexagons.update(line_path_hexs)
        for hex_id in edge_hexagons:
            hex_community_votes[hex_id].append(u_comm)

    final_hex_to_community = {}
    for hex_id, votes in hex_community_votes.items():
        if votes:
            most_common_community = Counter(votes).most_common(1)[0][0]
            final_hex_to_community[hex_id] = most_common_community
            
    hex_polygons = []
    for hex_id, community_id in final_hex_to_community.items():
        boundary = h3.h3_to_geo_boundary(hex_id, geo_json=True)
        hex_polygons.append({'geometry': Polygon(boundary), 'community_id': community_id})
    if not hex_polygons:
        print("⚠️ 未能生成任何H3社区边界，请检查输入数据或参数。")
        return gpd.GeoDataFrame()
        
    hex_gdf = gpd.GeoDataFrame(hex_polygons, crs="EPSG:4326")
    community_boundaries_gdf = hex_gdf.dissolve(by='community_id').reset_index()
    print(f"✅ H3边界描绘完成，最终生成 {len(community_boundaries_gdf)} 个社区多边形。")
    return community_boundaries_gdf

def calculate_community_features(communities_gdf, road_gdf, poi_gdf, G):
    # ... (此函数无需修改) ...
    print("\n--- 阶段4: 正在为每个社区计算分类特征 ---")
    communities_proj = communities_gdf.to_crs(epsg=3857) 
    road_proj = road_gdf.to_crs(epsg=3857)
    poi_proj = poi_gdf.to_crs(epsg=3857)
    communities_gdf['area_sqkm'] = communities_proj.area / 1_000_000
    road_clipped = gpd.sjoin(road_proj, communities_proj, how="inner", predicate="intersects")
    road_len_by_comm = road_clipped.groupby('community_id')['geometry'].apply(lambda geoms: geoms.length.sum())
    communities_gdf['road_density'] = communities_gdf['community_id'].map(road_len_by_comm) / (communities_gdf['area_sqkm'] * 1000)
    communities_gdf['road_density'].fillna(0, inplace=True)
    poi_clipped = gpd.sjoin(poi_proj, communities_proj, how="inner", predicate="within")
    poi_counts = poi_clipped.groupby(['community_id', 'category']).size().unstack(fill_value=0)
    communities_gdf = communities_gdf.merge(poi_counts, on='community_id', how='left').fillna(0)
    for col in poi_counts.columns:
        communities_gdf[f'poi_{col}_density'] = communities_gdf[col] / communities_gdf['area_sqkm']
        communities_gdf[f'poi_{col}_density'].fillna(0, inplace=True)
    print("  - 正在计算网络中心性...")
    centrality = nx.closeness_centrality(G)
    node_df = pd.DataFrame.from_dict({node: data for node, data in G.nodes(data=True)}, orient='index')
    node_df['centrality'] = node_df.index.map(centrality)
    avg_centrality = node_df.groupby('community')['centrality'].mean()
    communities_gdf['avg_centrality'] = communities_gdf['community_id'].map(avg_centrality).fillna(0)
    print("✅ 特征计算完成。")
    return communities_gdf

def classify_and_interpret_clusters(communities_gdf, k):
    # ... (此函数无需修改) ...
    print("\n--- 阶段5: 正在使用K-Means对社区进行分类 ---")
    feature_cols = ['area_sqkm', 'road_density', 'avg_centrality'] + [col for col in communities_gdf.columns if 'poi_' in col and '_density' in col]
    communities_gdf.replace([pd.NA, None, float('inf'), -float('inf')], 0, inplace=True) # 数据清洗
    X = communities_gdf[feature_cols].values
    X_scaled = StandardScaler().fit_transform(X)
    kmeans = KMeans(n_clusters=k, random_state=42, n_init='auto')
    communities_gdf['cluster_id'] = kmeans.fit_predict(X_scaled)
    print("\n--- 聚类结果解读 (各类别特征均值) ---")
    cluster_summary = communities_gdf.groupby('cluster_id')[feature_cols].mean()
    print(cluster_summary)
    print("\n--- 尝试自动命名聚类 (示例) ---")
    cluster_names = {}
    core_id = cluster_summary['avg_centrality'].idxmax()
    cluster_names[core_id] = "Urban Core / Commercial"
    rural_id = cluster_summary.sort_values(by=['area_sqkm', 'road_density'], ascending=[False, True]).index[0]
    if rural_id not in cluster_names: cluster_names[rural_id] = "Rural / Suburban"
    if 'poi_RESIDENTIAL_density' in cluster_summary.columns:
        residential_id = cluster_summary['poi_RESIDENTIAL_density'].idxmax()
        if residential_id not in cluster_names: cluster_names[residential_id] = "Residential Zone"
    remaining_ids = set(range(k)) - set(cluster_names.keys())
    for i, rid in enumerate(remaining_ids):
        cluster_names[rid] = f"Mixed Zone {i+1}"
    print(f"  - 命名结果: {cluster_names}")
    communities_gdf['cluster_name'] = communities_gdf['cluster_id'].map(cluster_names)
    return communities_gdf, cluster_names

def create_final_map(community_gdf, stores_gdf, boundary_gdf, cluster_names, output_path):
    # ... (此函数无需修改) ...
    print("\n--- 阶段6: 正在生成最终的分类地图 ---")
    map_center = [boundary_gdf.union_all().centroid.y, boundary_gdf.union_all().centroid.x]
    m = folium.Map(location=map_center, zoom_start=13, tiles="CartoDB positron")
    color_palette = {"Urban Core / Commercial": "#FF0000", "Residential Zone": "#0000FF", "Rural / Suburban": "#FFFF00", "Mixed Zone 1": "#008000"}
    cluster_id_to_name = dict(zip(cluster_names.keys(), cluster_names.values()))
    folium.GeoJson(
        community_gdf, name="社区功能分类 (H3边界)",
        style_function=lambda feature: {'fillColor': color_palette.get(cluster_id_to_name.get(feature['properties']['cluster_id']), '#808080'), 'color': 'black', 'weight': 0.5, 'fillOpacity': 0.7},
        tooltip=folium.GeoJsonTooltip(fields=['cluster_name', 'community_id'], aliases=['功能区类型:', '社区ID:'])
    ).add_to(m)
    folium.GeoJson(boundary_gdf, name="研究区域边界", style_function=lambda x: {'color': 'black', 'weight': 2.5, 'fillOpacity': 0}).add_to(m)
    stores_group = folium.FeatureGroup(name="现有瑞幸门店")
    for _, store in stores_gdf.iterrows():
        folium.Marker(location=[store.geometry.y, store.geometry.x], icon=folium.Icon(color='darkgreen', icon='coffee', prefix='fa')).add_to(stores_group)
    stores_group.add_to(m)
    folium.LayerControl().add_to(m)
    m.save(output_path)
    print(f"✅ 地图已成功保存至: {output_path}")

def main():
    road_gdf, poi_gdf, stores_gdf, boundary_gdf = load_data(BOUNDARY_SHP, ROAD_SHP, POI_CSV, STORES_CSV)
    if road_gdf is None: return

    road_graph = perform_louvain_community_detection(road_gdf, LOUVAIN_RESOLUTION)
    # [核心] 调用H3描边函数，而不是凸包函数
    h3_boundaries = delineate_communities_with_h3(road_graph, H3_RESOLUTION, MIN_COMMUNITY_SIZE)
    
    if h3_boundaries.empty: return

    communities_with_features = calculate_community_features(h3_boundaries, road_gdf, poi_gdf, road_graph)
    classified_communities, cluster_names = classify_and_interpret_clusters(communities_with_features, NUM_CLUSTERS)
    create_final_map(classified_communities, stores_gdf, boundary_gdf, cluster_names, OUTPUT_MAP_HTML)

if __name__ == "__main__":
    main()